import openpyxl
from openpyxl import Workbook
import os

# SAVE_DIR = ".\\WeChatTextExport"  # 导出目录，可自行修改
# EXCEL_NAME = "mapping.xlsx"

# # 创建保存目录
# if not os.path.exists(SAVE_DIR):
#     os.makedirs(SAVE_DIR)

class Wxname2StudentID:
    """
    Excel 映射管理类：
    功能：
      1. 从 Excel 加载 {微信号 → 学号} 映射
      2. 将映射写回 Excel（自动去重、合并、排序）
    """

    def __init__(self, excel_path):
        self.excel_path = excel_path

    # ======================================================
    # 1. 读取 Excel 映射（文件不存在 → 返回 {}）
    # ======================================================
    def load_mapping(self, sheet_name="Sheet1"):
        """
        读取 Excel 映射，返回 dict {wx_id: student_id}
        文件不存在则返回空 map，不报错
        """

        try:
            wb = openpyxl.load_workbook(self.excel_path)
        except FileNotFoundError:
            return {}  # 文件不存在 → 空字典

        if sheet_name not in wb.sheetnames:
            return {}  # sheet 不存在 → 空字典

        sheet = wb[sheet_name]

        mapping = {}

        for row in sheet.iter_rows(min_row=2, values_only=True):
            wx_id, student_id = row[0], row[1]
            if wx_id and student_id:
                mapping[str(wx_id).strip()] = str(student_id).strip()

        return mapping

    # ======================================================
    # 2. 写回 Excel（自动去重、更新、排序）
    # ======================================================
    def save_mapping(self, mapping, sheet_name="Sheet1"):
        """
        保存映射到 Excel：
        - 微信号不能重复：重复只更新，不新增
        - 写入后按“学号”排序（升序）
        """

        try:
            wb = openpyxl.load_workbook(self.excel_path)
            sheet = wb[sheet_name] if sheet_name in wb.sheetnames else wb.create_sheet(sheet_name)
        except FileNotFoundError:
            wb = Workbook()
            sheet = wb.active
            sheet.title = sheet_name

        # 如果空表 → 写表头
        if sheet.max_row == 1 and not sheet["A1"].value:
            sheet["A1"] = "微信号"
            sheet["B1"] = "学号"

        # 先把 Excel 里的数据读出来，构成 existing_dict
        existing_dict = {}
        for row in sheet.iter_rows(min_row=2, values_only=True):
            wx_id, student_id = row[0], row[1]
            if wx_id and student_id:
                existing_dict[str(wx_id).strip()] = str(student_id).strip()

        # 合并（mapping 更新 existing_dict）
        for wx_id in mapping:
            new_value = mapping.get(wx_id)
            if new_value:
                existing_dict[str(wx_id).strip()] = str(new_value).strip()

        # 排序（按学号升序）
        sorted_items = sorted(
            existing_dict.items(),
            key=lambda x: int(x[1]) if x[1].isdigit() else x[1]
        )

        # 清空并重写
        sheet.delete_rows(2, sheet.max_row)

        row_index = 2
        for wx_id, student_id in sorted_items:
            sheet[f"A{row_index}"] = wx_id
            sheet[f"B{row_index}"] = student_id
            row_index += 1

        wb.save(self.excel_path)

# def main():
#     a = Wxname2StudentID(os.path.join(SAVE_DIR, EXCEL_NAME))
#     mapping = a.load_mapping()
#     print(mapping)
#     mapping['adddd'] = 29180
#     a.save_mapping(mapping)

# if __name__ == "__main__":
#     main()